"""
数据管理模块
负责股票数据的获取、存储和管理
"""

import os
import json
import time
import logging
from typing import Dict, List, Any, Optional
from datetime import datetime, timedelta
import pandas as pd

logger = logging.getLogger(__name__)

class DataManager:
    """数据管理器"""
    
    def __init__(self):
        """初始化数据管理器"""
        self.data_dir = "data"
        self.cache_dir = os.path.join(self.data_dir, "cache")
        self.stock_list_file = os.path.join(self.data_dir, "stock_list.json")
        self.config_file = os.path.join(self.data_dir, "data_config.json")
        
        # 确保目录存在
        os.makedirs(self.data_dir, exist_ok=True)
        os.makedirs(self.cache_dir, exist_ok=True)
        
        # 加载配置
        self.config = self._load_config()
    
    def _load_config(self) -> Dict[str, Any]:
        """加载配置"""
        default_config = {
            "data_sources": {
                "primary": "eastmoney",
                "backup": ["pytdx", "akshare"]
            },
            "update_interval": 3600,  # 1小时
            "cache_expiry": 86400,    # 24小时
            "max_retries": 3,
            "timeout": 30
        }
        
        if os.path.exists(self.config_file):
            try:
                with open(self.config_file, 'r', encoding='utf-8') as f:
                    config = json.load(f)
                    # 合并默认配置
                    for key, value in default_config.items():
                        if key not in config:
                            config[key] = value
                    return config
            except Exception as e:
                logger.error(f"加载配置文件失败: {e}")
        
        return default_config
    
    def _save_config(self):
        """保存配置"""
        try:
            with open(self.config_file, 'w', encoding='utf-8') as f:
                json.dump(self.config, f, ensure_ascii=False, indent=2)
        except Exception as e:
            logger.error(f"保存配置文件失败: {e}")
    
    def is_available(self) -> bool:
        """检查数据管理器是否可用"""
        try:
            # 检查必要的依赖
            import requests
            import pandas as pd
            return True
        except ImportError:
            return False
    
    def check_data_sources(self) -> Dict[str, Dict[str, Any]]:
        """检查数据源状态"""
        sources = {}
        
        # 检查东方财富
        sources['东方财富'] = self._check_eastmoney()
        
        # 检查pytdx
        sources['pytdx'] = self._check_pytdx()
        
        # 检查akshare
        sources['akshare'] = self._check_akshare()
        
        return sources
    
    def _check_eastmoney(self) -> Dict[str, Any]:
        """检查东方财富数据源"""
        try:
            import requests
            session = requests.Session()
            session.trust_env = False
            
            url = "http://push2.eastmoney.com/api/qt/clist/get"
            params = {'pn': '1', 'pz': '1'}
            
            start_time = time.time()
            response = session.get(url, params=params, timeout=10)
            response_time = (time.time() - start_time) * 1000
            
            if response.status_code == 200:
                return {
                    'connected': True,
                    'response_time': response_time,
                    'status': '连接正常'
                }
            else:
                return {
                    'connected': False,
                    'response_time': None,
                    'status': f'HTTP {response.status_code}'
                }
        except Exception as e:
            return {
                'connected': False,
                'response_time': None,
                'status': f'连接失败: {str(e)}'
            }
    
    def _check_pytdx(self) -> Dict[str, Any]:
        """检查pytdx数据源"""
        try:
            from pytdx.hq import TdxHq_API
            api = TdxHq_API()
            servers = [('221.131.136.194', 7709), ('111.229.200.45', 7709)]
            
            for host, port in servers:
                try:
                    start_time = time.time()
                    if api.connect(host, port):
                        response_time = (time.time() - start_time) * 1000
                        api.disconnect()
                        return {
                            'connected': True,
                            'response_time': response_time,
                            'status': '连接正常'
                        }
                except:
                    continue
            
            return {
                'connected': False,
                'response_time': None,
                'status': '所有服务器连接失败'
            }
        except ImportError:
            return {
                'connected': False,
                'response_time': None,
                'status': '模块未安装'
            }
        except Exception as e:
            return {
                'connected': False,
                'response_time': None,
                'status': f'连接失败: {str(e)}'
            }
    
    def _check_akshare(self) -> Dict[str, Any]:
        """检查akshare数据源"""
        try:
            import akshare as ak
            start_time = time.time()
            df = ak.tool_trade_date_hist_sina()
            response_time = (time.time() - start_time) * 1000
            
            if df is not None and len(df) > 0:
                return {
                    'connected': True,
                    'response_time': response_time,
                    'status': '连接正常'
                }
            else:
                return {
                    'connected': False,
                    'response_time': None,
                    'status': '返回数据为空'
                }
        except ImportError:
            return {
                'connected': False,
                'response_time': None,
                'status': '模块未安装'
            }
        except Exception as e:
            return {
                'connected': False,
                'response_time': None,
                'status': f'连接失败: {str(e)}'
            }
    
    def update_stock_list(self):
        """更新股票列表"""
        print("正在更新股票列表...")
        
        try:
            # 尝试从东方财富获取股票列表
            stocks = self._get_stocks_from_eastmoney()
            
            if stocks:
                # 保存股票列表
                with open(self.stock_list_file, 'w', encoding='utf-8') as f:
                    json.dump(stocks, f, ensure_ascii=False, indent=2)
                
                print(f"成功更新股票列表，共 {len(stocks)} 只股票")
                logger.info(f"更新股票列表成功，共 {len(stocks)} 只股票")
            else:
                print("获取股票列表失败")
                logger.error("获取股票列表失败")
                
        except Exception as e:
            print(f"更新股票列表时发生错误: {e}")
            logger.error(f"更新股票列表失败: {e}")
    
    def _get_stocks_from_eastmoney(self) -> List[Dict[str, str]]:
        """从东方财富获取股票列表"""
        try:
            import requests
            
            session = requests.Session()
            session.trust_env = False
            
            url = "http://push2.eastmoney.com/api/qt/clist/get"
            params = {
                'pn': 1,
                'pz': 5000,
                'fs': 'm:0+t:6,m:0+t:13,m:0+t:80,m:1+t:2,m:1+t:23',
                'fields': 'f12,f14'
            }
            
            response = session.get(url, params=params, timeout=30)
            
            if response.status_code == 200:
                data = response.json()
                if 'data' in data and 'diff' in data['data']:
                    stocks = []
                    for item in data['data']['diff']:
                        if item.get('f12') and item.get('f14'):
                            stocks.append({
                                'code': item['f12'],
                                'name': item['f14']
                            })
                    return stocks
            
            return []
            
        except Exception as e:
            logger.error(f"从东方财富获取股票列表失败: {e}")
            return []
    
    def download_historical_data(self):
        """下载历史数据"""
        print("正在下载历史数据...")
        
        try:
            # 加载股票列表
            if not os.path.exists(self.stock_list_file):
                print("股票列表不存在，请先更新股票列表")
                return
            
            with open(self.stock_list_file, 'r', encoding='utf-8') as f:
                stocks = json.load(f)
            
            # 限制下载数量以避免超时
            stocks_to_download = stocks[:100]  # 只下载前100只股票
            
            print(f"开始下载 {len(stocks_to_download)} 只股票的历史数据...")
            
            success_count = 0
            for i, stock in enumerate(stocks_to_download):
                try:
                    stock_code = stock['code']
                    stock_name = stock['name']
                    
                    # 下载股票数据
                    data = self._download_stock_data(stock_code)
                    
                    if data is not None:
                        # 保存数据
                        filename = os.path.join(self.cache_dir, f"{stock_code}.csv")
                        data.to_csv(filename, index=False, encoding='utf-8')
                        success_count += 1
                    
                    # 显示进度
                    if (i + 1) % 10 == 0:
                        print(f"已处理 {i + 1}/{len(stocks_to_download)} 只股票")
                    
                    # 避免请求过于频繁
                    time.sleep(0.1)
                    
                except Exception as e:
                    logger.error(f"下载股票 {stock['code']} 数据失败: {e}")
                    continue
            
            print(f"历史数据下载完成，成功下载 {success_count} 只股票的数据")
            
        except Exception as e:
            print(f"下载历史数据时发生错误: {e}")
            logger.error(f"下载历史数据失败: {e}")
    
    def _download_stock_data(self, stock_code: str) -> Optional[pd.DataFrame]:
        """下载单只股票的历史数据"""
        try:
            import akshare as ak
            
            # 获取股票历史数据
            df = ak.stock_zh_a_hist(symbol=stock_code, period="daily", 
                                  start_date="20200101", end_date="20241231")
            
            if df is not None and not df.empty:
                return df
            else:
                return None
                
        except Exception as e:
            logger.error(f"下载股票 {stock_code} 数据失败: {e}")
            return None
    
    def check_data_quality(self):
        """检查数据质量"""
        print("正在检查数据质量...")
        
        try:
            # 检查股票列表
            if os.path.exists(self.stock_list_file):
                with open(self.stock_list_file, 'r', encoding='utf-8') as f:
                    stocks = json.load(f)
                print(f"股票列表: {len(stocks)} 只股票")
            else:
                print("股票列表文件不存在")
            
            # 检查历史数据文件
            if os.path.exists(self.cache_dir):
                data_files = [f for f in os.listdir(self.cache_dir) if f.endswith('.csv')]
                print(f"历史数据文件: {len(data_files)} 个")
                
                # 检查数据完整性
                valid_files = 0
                for file in data_files:
                    file_path = os.path.join(self.cache_dir, file)
                    try:
                        df = pd.read_csv(file_path)
                        if len(df) > 0:
                            valid_files += 1
                    except:
                        continue
                
                print(f"有效数据文件: {valid_files} 个")
            else:
                print("历史数据目录不存在")
                
        except Exception as e:
            print(f"检查数据质量时发生错误: {e}")
            logger.error(f"检查数据质量失败: {e}")
    
    def clean_cache(self):
        """清理缓存数据"""
        print("正在清理缓存数据...")
        
        try:
            if os.path.exists(self.cache_dir):
                files = os.listdir(self.cache_dir)
                for file in files:
                    file_path = os.path.join(self.cache_dir, file)
                    if os.path.isfile(file_path):
                        os.remove(file_path)
                
                print(f"已清理 {len(files)} 个缓存文件")
            else:
                print("缓存目录不存在")
                
        except Exception as e:
            print(f"清理缓存时发生错误: {e}")
            logger.error(f"清理缓存失败: {e}")
    
    def get_stock_list(self) -> List[Dict[str, str]]:
        """获取股票列表"""
        if os.path.exists(self.stock_list_file):
            try:
                with open(self.stock_list_file, 'r', encoding='utf-8') as f:
                    return json.load(f)
            except Exception as e:
                logger.error(f"读取股票列表失败: {e}")
        
        return []
    
    def get_stock_data(self, stock_code: str) -> Optional[pd.DataFrame]:
        """获取股票历史数据"""
        filename = os.path.join(self.cache_dir, f"{stock_code}.csv")
        
        if os.path.exists(filename):
            try:
                return pd.read_csv(filename)
            except Exception as e:
                logger.error(f"读取股票 {stock_code} 数据失败: {e}")
        
        return None
